// Copyright (c) 2022. China Mobile(SuZhou)Software Technology Co.,Ltd. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package server

import (
	"flag"
	"io/ioutil"

	"gitee.com/he3db/he3proxy/core/errors"
	"gitee.com/he3db/he3proxy/core/golog"
	"gitee.com/he3db/he3proxy/postgresql/hba"
)

// define hba configuration
var HbaConfigFile = flag.String("hba", "./etc/hba.conf", "He3Proxy hba config file")

var (
	hbaAuthMethods = map[string]methodInfo{}
)

type methodInfo struct {
	validConnTypes hba.ConnType
}

// RegisterAuthMethod registers an AuthMethod for pgwire
// authentication and for use in HBA configuration.
//
// The minReqVersion is checked upon configuration to verify whether
// the current active cluster version is at least the version
// specified.
//
// The validConnTypes is checked during rule matching when accepting
// connections: if the connection type is not accepted by the auth
// method, authentication is refused upfront. For example, the "cert"
// method requires SSL; if a rule specifies "host .... cert" and the
// client connects without SSL, the authentication is refused.
// (To express "cert on SSL, password on non-SSL", the HBA conf
// can list 'hostssl ... cert; hostnossl .... password' instead.)
//
// The checkEntry method, if provided, is called upon configuration
// the cluster setting in the SQL client which attempts to change the
// configuration. It can block the configuration if e.g. the syntax is
// invalid.
func RegisterAuthMethod(
	method string, validConnTypes hba.ConnType,
) {
	hbaAuthMethods[method] = methodInfo{validConnTypes}
}

func loadDefaultMethods() {
	// The "password" method requires a clear text password.
	//
	// Care should be taken by administrators to only accept this auth
	// method over secure connections, e.g. those encrypted using SSL.
	RegisterAuthMethod("md5", hba.ConnAny)

	// The "reject" method rejects any connection attempt that matches
	// the current rule.
	RegisterAuthMethod("reject", hba.ConnAny)

	// The "trust" method accepts any connection attempt that matches
	// the current rule.
	RegisterAuthMethod("trust", hba.ConnAny)
}

// DefaultHBAConfig is used when the stored HBA configuration string
// is empty or invalid.
var DefaultHBAConfig = func() *hba.Conf {
	loadDefaultMethods()
	conf, err := ParseAndNormalize(`
host  all all  all md5 # built-in CockroachDB default
local all all      md5      # built-in CockroachDB default
`)
	if err != nil {
		panic(err)
	}
	return conf
}()

// GetAuthenticationConfiguration retrieves the current applicable
// authentication configuration.
//
// This is guaranteed to return a valid configuration. Additionally,
// the various setters for the configuration also pass through
// ParseAndNormalize(), whereby an entry is always present at the start,
// to enable root to log in with a valid client cert.
//
// The data returned by this method is also observable via the debug
// endpoint /debug/hba_conf.
func GetAuthenticationConfiguration() *hba.Conf {
	data, err := ioutil.ReadFile(*HbaConfigFile)
	if err != nil {
		golog.Error("hba", "GetAuthenticationConfiguration", err.Error(), 0)
		return DefaultHBAConfig
	}

	auth, err := ParseAndNormalize(string(data))
	if auth == nil {
		golog.Error("hba_conf", "GetAuthenticationConfiguration", err.Error(), 0)
		golog.Warn("hba_conf", "GetAuthenticationConfiguration", "Will use default configuration for no authentication", 0)
		// This can happen when using the value for the first time before
		// the cluster setting has ever been set.
		auth = DefaultHBAConfig
	}
	return auth
}

// ParseAndNormalize calls hba.ParseAndNormalize and also ensures the
// configuration starts with a rule that authenticates the root user
// with client certificates.
//
// This prevents users from shooting themselves in the foot and making
// root not able to login, thus disallowing anyone from fixing the HBA
// configuration.
func ParseAndNormalize(val string) (*hba.Conf, error) {
	conf, err := hba.ParseAndNormalize(val)
	if err != nil {
		return conf, err
	}

	if len(conf.Entries) == 0 || !conf.Entries[0].Equivalent(rootEntry) {
		entries := make([]hba.Entry, 1, len(conf.Entries)+1)
		entries[0] = rootEntry
		entries = append(entries, conf.Entries...)
		conf.Entries = entries
	}

	// Lookup and cache the auth methods.
	for i := range conf.Entries {
		method := conf.Entries[i].Method.Value
		info, ok := hbaAuthMethods[method]
		if !ok {
			// TODO(knz): Determine if an error should be reported
			// upon unknown auth methods.
			// See: https://github.com/cockroachdb/cockroach/issues/43716
			return nil, errors.ErrFormat("unknown auth method %s", method)
		}
		conf.Entries[i].MethodFn = info
	}

	return conf, nil
}

var insecureEntry = hba.Entry{
	ConnType: hba.ConnHostAny,
	User:     []hba.String{{Value: "all", Quoted: false}},
	Address:  hba.AnyAddr{},
	Method:   hba.String{Value: "--insecure"},
}

var rootEntry = hba.Entry{
	ConnType: hba.ConnHostAny,
	User:     []hba.String{{Value: "root", Quoted: false}},
	Address:  hba.AnyAddr{},
	Method:   hba.String{Value: "md5"},
	Input:    "host  all root all md5 # CockroachDB mandatory rule",
}
